import cv2
import matplotlib.pyplot as plt
import numpy as np
import ezdxf
doc = ezdxf.new(dxfversion='R2010')  # 创建新 DXF 文档
msp = doc.modelspace()  # 获取模型空间
# 读取图像
image = cv2.imread('./data/aa.png')

# 原始图像副本
original_image = image.copy()

# 转换为灰度图像
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
# 调整后的参数
edges = cv2.Canny(gray,80, 200, apertureSize=3)
lines = cv2.HoughLinesP(edges, rho=1, theta=np.pi*3 / 180, threshold=25, minLineLength=25, maxLineGap=30)

# 创建一个空白图像用于绘制线段（与原始图像大小相同）
line_image = np.zeros_like(original_image)

# 绘制检测到的线段
if lines is not None:
    for line in lines:
        x1, y1, x2, y2 = line[0]
        # 生成较亮的随机颜色（避免深色）
        color = tuple(np.random.randint(150, 255, 3).tolist())
        # 在空白图像上绘制线段
        cv2.line(line_image, (x1, y1), (x2, y2), color, 2)
        # 绘制起点
        cv2.circle(line_image, (x1, y1), 5, (0, 0, 255), -1)  # 红色实心圆表示起点
        # 绘制终点
        cv2.circle(line_image, (x2, y2), 5, (255, 0, 0), -1)  # 蓝色实心圆表示终点

        # 添加线段
        msp.add_line((x1, y1), (x2, y2), dxfattribs={'color': 7})  # 默认颜色为白色
        
        # 添加起点圆
        msp.add_circle((x1, y1), radius=5, dxfattribs={'color': 1})  # 红色表示起点
        
        # 添加终点圆
        msp.add_circle((x2, y2), radius=5, dxfattribs={'color': 5})  # 蓝色表示终点
plt.figure(figsize=(18, 6))
doc.saveas('output/detected_lines.dxf')  # 保存 DXF 文件
# 原始图像
plt.subplot(1, 3, 1)
plt.imshow(cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB))
plt.title("Original Image")
plt.axis('off')

# 边缘检测图像
plt.subplot(1, 3, 2)
plt.imshow(edges, cmap='gray')  # 显示 Canny 边缘图像
plt.title("Canny Edges")
plt.axis('off')

# 线段检测结果（无原图背景）
plt.subplot(1, 3, 3)
plt.imshow(cv2.cvtColor(line_image, cv2.COLOR_BGR2RGB))
plt.title("Detected Lines (No Background)")
plt.axis('off')

plt.tight_layout()
plt.show()